#ifndef RAY_TRACE_PRO_BVH_H
#define RAY_TRACE_PRO_BVH_H

#include "Basic/GxLibBasic.h"
#include "RayTraceSystem.h"
#include "Debug.h"
#include "BBox.h"
#include "Basic/Pool.hpp"
#include <ppl.h>
#include "Windows/GxThreads.h"
#include "Windows/GxDebug.h"

using namespace Concurrency;

using namespace GxLibBasic;

#define PROFILE_BVH

namespace RayTracePro
{
	const int nBuckets = 64;

	class BvhNode
	{
	public:
		BBox Bounds;
		unsigned int Axis:2;
#ifdef PROFILE_BVH
		unsigned int IsBadNode:1;
		int ElementCount:29;
#else
		int ElementCount:30;
#endif
		union
		{
			int ElementId;
			int ChildOffset;
		};
		inline bool GetIsLeaf()
		{
			return ElementCount != 0;
		}
		inline int GetElementCount()
		{
			return ElementCount;
		}
	};

	template<typename T>
	class BvhNode_Build
	{
	public:
		BBox Bounds;
		int Axis;
		T** Elements;
		int ElementCount;
		BvhNode_Build* Children[2];
#ifdef PROFILE_BVH
		bool IsBadNode;
#endif
		void AllocElements(int count)
		{
			Elements = (T**)Alloc(count*sizeof(T*));
			ElementCount = count;
		}
		void FreeElements()
		{
			if (Elements)
			{
				Free(Elements);
				Elements = 0;
			}
		}
		BvhNode_Build()
		{
			Children[0] = 0;
			Children[1] = 0;
			Axis = 0;
			ElementCount = 0;
			Elements = 0;
			IsBadNode = false;
		}
		~BvhNode_Build()
		{
			if (Children[0])
				delete Children[0];
			if (Children[1])
				delete Children[1];
			FreeElements();
		}
		static void* operator new(size_t size)
		{
			return Alloc(size);
		}
		static void operator delete(void *p)
		{
			return Free(p);
		}

	};

	template<typename T>
	class Bvh_Build
	{
	public:
		RefPtr<BvhNode_Build<T>> Root;
		IntPtr ElementListSize;
		IntPtr NodeCount;
#ifdef PROFILE_BVH
		volatile long long Profile_BadSplits;
		GxWinSystem::Threading::SpinLock Profile_Lock;

#endif
		Bvh_Build()
		{
			NodeCount = 0;
			
		}
	};

	template<typename T>
	class Bvh
	{
	private:
		int FlattenNodes(BvhNode_Build<T> * node)
		{
			int id = Nodes.Count();
			BvhNode n;
			n.Axis = node->Axis;
			n.Bounds = node->Bounds;
			if (node->Elements == 0)
				n.ElementCount = 0;
			else
				n.ElementCount = node->ElementCount;
#ifdef PROFILE_BVH
			n.IsBadNode = node->IsBadNode;
#endif
			Nodes.Add(n);
			
			if (node->Elements == 0)
			{
				FlattenNodes(node->Children[0]);
				Nodes[id].ChildOffset = FlattenNodes(node->Children[1]) - id;
			}
			else
			{
				Nodes[id].ElementId = Elements.Count();
				for (int i = 0; i<node->ElementCount; i++)
					Elements.Add(*node->Elements[i]);
			}
			return id;
		}
	public:
		Array<BvhNode, true> Nodes;
		Array<T, true> Elements;
		void FromBuild(Bvh_Build<T> &bvh)
		{
			Nodes.Clear();
			Elements.Clear();
			Nodes.SetCapacity((int)bvh.NodeCount);
			Elements.SetCapacity((int)bvh.ElementListSize);
			FlattenNodes(bvh.Root.operator->());
		}
		~Bvh()
		{}
	};

	template<typename T>
	class BuildData
	{
	public:
		T * Element;
		BBox Bounds;
		Vec3 Center;
	};

	inline void Assert(bool pred)
	{
#ifdef _DEBUG
		if (!pred)
			throw 0;
#endif
	}

	inline float SurfaceArea(BBox & box)
	{
		return ((box.xMax-box.xMin)*(box.yMax-box.yMin) + (box.xMax-box.xMin)*(box.zMax-box.zMax) + (box.yMax-box.yMin)*(box.zMax-box.zMin))*2.0f;
	}

	struct BucketInfo
	{
		int count;
		BBox bounds;
		BucketInfo()
		{
			count = 0;
			bounds.Init();
		}
	};

	template<typename T, typename CostEvaluator>
	BvhNode_Build<T> * ConstructBvhNode(Bvh_Build<T> & tree, BuildData<T>* elements, IntPtr elementCount, IntPtr & elementListSize, IntPtr & nodeCount, CostEvaluator & eval, int depth)
	{
		BvhNode_Build<T> * node = new BvhNode_Build<T>();
		//_InterlockedIncrement((volatile long*)(&tree.NodeCount));
		nodeCount = 1;
		elementListSize = 0;
		if (elementCount == 1 || depth == 61)
		{
			node->Bounds = elements->Bounds;
			node->AllocElements(1);
			node->Elements[0] = elements->Element;
			elementListSize = 1;
			return node;
		}
		else
		{
			BBox centroidBounds;
			BBox bbox;
			centroidBounds.Init();
			bbox.Init();
			for (int i = 0; i < elementCount; i++)
			{
				centroidBounds.Union(elements[i].Center);
				bbox.Union(elements[i].Bounds);
			}
			node->Bounds = bbox;
			int dim = centroidBounds.MaxDimension();

			if (centroidBounds.Min[dim] == centroidBounds.Max[dim])
			{
				node->Bounds = bbox;
				node->AllocElements((int)elementCount);
				for (int i = 0; i < (int)elementCount; i++)
				{
					node->Elements[i] = elements[i].Element;
				}
				elementListSize = elementCount;
				return node;
			}

			BucketInfo buckets[nBuckets];
			if (elementCount > (2<<12))
			{
				const int processorCount = 16;
				BucketInfo buckets_proc[processorCount][nBuckets];
				int blockSize = (int)(elementCount/processorCount);
				parallel_for(0, processorCount, [&](int procId)
				{
					int end;
					if (procId == processorCount - 1)
						end = (int)elementCount;
					else
						end = (procId+1)*blockSize;
					for (int i = procId*blockSize; i<end; i++)
					{
						int b = (int)(nBuckets *
							((elements[i].Center[dim] - centroidBounds.Min[dim]) / (centroidBounds.Max[dim] - centroidBounds.Min[dim])));
						if (b == nBuckets) b = nBuckets-1;
						Assert(b >= 0 && b < nBuckets);
						buckets_proc[procId][b].count++;
						buckets_proc[procId][b].bounds.Union(elements[i].Bounds);
					}
				});
				for (int i = 0; i<nBuckets; i++)
				{
					for (int j = 0; j<processorCount; j++)
					{
						buckets[i].count += buckets_proc[j][i].count;
						buckets[i].bounds.Union(buckets_proc[j][i].bounds);
					}
				}
			}
			else
			{
				for (int i = 0; i<elementCount; i++)
				{
					int b = (int)(nBuckets *
						((elements[i].Center[dim] - centroidBounds.Min[dim]) / (centroidBounds.Max[dim] - centroidBounds.Min[dim])));
					if (b == nBuckets) b = nBuckets-1;
					Assert(b >= 0 && b < nBuckets);
					buckets[b].count++;
					buckets[b].bounds.Union(elements[i].Bounds);
				}
			}

			float cost[nBuckets-1];
            for (int i = 0; i < nBuckets-1; i++)
			{
                BBox b0, b1;
				b0.Init();
				b1.Init();
                int count0 = 0, count1 = 0;
                for (int j = 0; j <= i; j++)
				{
                    b0.Union(buckets[j].bounds);
                    count0 += buckets[j].count;
                }
                for (int j = i+1; j < nBuckets; j++)
				{
                    b1.Union(buckets[j].bounds);
                    count1 += buckets[j].count;
                }

                cost[i] = eval.EvalCost(count0, SurfaceArea(b0), count1, SurfaceArea(b1), SurfaceArea(bbox));
            }
 
            float minCost = cost[0];
            int minCostSplit = 0;
            for (int i = 1; i < nBuckets-1; i++)
			{
                if (cost[i] < minCost)
				{
                    minCost = cost[i];
                    minCostSplit = i;
                }
            }

            if (elementCount > CostEvaluator::ElementsPerNode ||
                minCost < elementCount)
			{
                BuildData<T> *pmid = std::partition(elements,
                    elements + elementCount,
					[&](const BuildData<T> &p)
					{
						int b = (int)(nBuckets * ((p.Center[dim] - centroidBounds.Min[dim]) /
							 (centroidBounds.Max[dim] - centroidBounds.Min[dim])));
						if (b == nBuckets) b = nBuckets-1;
						Assert(b >= 0 && b < nBuckets);
						return b <= minCostSplit;
					});
				node->Axis = dim;
				IntPtr listSize1, listSize2;
				IntPtr nodeCount1, nodeCount2;
#ifdef PROFILE_BVH
				if (minCostSplit == 0 && depth<=32)
				{
					float endCost = eval.EvalCost(0, (bbox.Max[(dim+1)%3]-bbox.Min[(dim+1)%3])*(bbox.Max[(dim+2)%3]-bbox.Min[(dim+2)%3]), (int)elementCount, SurfaceArea(bbox), SurfaceArea(bbox));
					if (endCost < minCost)
					{
						int valAdd = 1;//1<<(32-depth);
						tree.Profile_Lock.Lock();
						tree.Profile_BadSplits += valAdd;
						tree.Profile_Lock.Unlock();
						node->IsBadNode = true;
					}
				}
				if (pmid == elements)
					pmid = elements + 1;
				else if (pmid == elements + elementCount)
					pmid = elements + elementCount-1;
#endif
				if (depth > 8)
				{
					node->Children[0] = ConstructBvhNodeNonRec<T, CostEvaluator>(tree, elements, pmid-elements, listSize1, nodeCount1, eval, depth+1);
					node->Children[1] = ConstructBvhNodeNonRec<T, CostEvaluator>(tree, pmid, elements + elementCount - pmid, listSize2, nodeCount2, eval, depth+1);
				}
				else
				{
					parallel_invoke(
						[&]()
						{
							node->Children[0] = ConstructBvhNode<T, CostEvaluator>(tree, elements, pmid-elements, listSize1, nodeCount1, eval, depth+1);
						},
						[&]()
						{
							node->Children[1] = ConstructBvhNode<T, CostEvaluator>(tree, pmid, elements + elementCount - pmid, listSize2, nodeCount2, eval, depth+1);
						});
				}
				node->ElementCount = (int)(elementListSize = listSize1+listSize2);
				nodeCount += nodeCount1+nodeCount2;
			}
            else
			{
				node->AllocElements((int)elementCount);
				node->Bounds = bbox;
				for (int i = 0; i < (int)elementCount; i++)
				{
					node->Elements[i] = elements[i].Element;
				}
				elementListSize = elementCount;
            }
			return node;
		}
	}

	template<typename T, typename CostEvaluator>
	BvhNode_Build<T> * ConstructBvhNodeNonRec(Bvh_Build<T> & tree, BuildData<T>* elements, IntPtr elementCount, IntPtr & elementListSize, IntPtr & nodeCount, CostEvaluator & eval, int depth)
	{
		Int64 Profile_BadSplits = 0;
		struct BvhJob
		{
			BvhNode_Build<T> ** result;
			BuildData<T>* elements;
			IntPtr elementCount;
			BvhJob()
			{}
			BvhJob(BvhNode_Build<T> ** result, BuildData<T>* elements, IntPtr elementCount)
			{
				this->elements = elements;
				this->elementCount = elementCount;
				this->result = result;
			}
		};
		const int stackSize = 256;
		BvhJob stack[stackSize];
		int stackPtr = 0;
		auto pushJob = [&](BvhNode_Build<T> ** result, BuildData<T>* elements, IntPtr elementCount)
		{
			BvhJob job(result, elements, elementCount);
			if (stackPtr < stackSize)
				stack[stackPtr++] = job;
			else
				throw "stack overflow";
		};
		auto popJob = [&]()->BvhJob
		{
			if (stackPtr)
				return stack[--stackPtr];
			else
				throw "stack empty";
		};

		BvhNode_Build<T> * rs = 0;
		nodeCount = 0;
		elementListSize = 0;
		BvhJob job(&rs, elements, elementCount);
		while (true)
		{
			BvhNode_Build<T> * node = new BvhNode_Build<T>();
			nodeCount++;
			(*job.result) = node;
			BuildData<T>* elements = job.elements;
			IntPtr elementCount = job.elementCount;
			if (elementCount == 0)
			{
				printf("elementCount = 0 !");
				throw 0;
			}
			if (elementCount == 1 || stackPtr == stackSize)
			{
				node->Bounds = elements->Bounds;
				node->AllocElements((int)elementCount);
				for (int i = 0; i < (int)elementCount; i++)
				{
					node->Elements[i] = elements[i].Element;
				}
				elementListSize += elementCount;
				if (!stackPtr)
					break;
				else
					job = popJob();
				continue;
			}
			else
			{
				BBox centroidBounds;
				BBox bbox;
				centroidBounds.Init();
				bbox.Init();
				for (int i = 0; i < elementCount; i++)
				{
					centroidBounds.Union(elements[i].Center);
					bbox.Union(elements[i].Bounds);
				}
				node->Bounds = bbox;
				int dim = centroidBounds.MaxDimension();

				if (centroidBounds.Min[dim] == centroidBounds.Max[dim])
				{
					node->Bounds = bbox;
					node->AllocElements((int)elementCount);
					for (int i = 0; i < (int)elementCount; i++)
					{
						node->Elements[i] = elements[i].Element;
					}
					elementListSize += elementCount;
					if (!stackPtr)
						break;
					else
						job = popJob();
					continue;
				}

				BucketInfo buckets[nBuckets];
				for (int i = 0; i<elementCount; i++)
				{
					int b = (int)(nBuckets *
						((elements[i].Center[dim] - centroidBounds.Min[dim]) / (centroidBounds.Max[dim] - centroidBounds.Min[dim])));
					if (b == nBuckets) b = nBuckets-1;
					Assert(b >= 0 && b < nBuckets);
					buckets[b].count++;
					buckets[b].bounds.Union(elements[i].Bounds);
				}
				
				float minCost = FLT_MAX;
				int minCostSplit = 0;
#ifdef PROFILE_BVH
				BBox minCostB1, minCostB2;
#endif
				for (int i = 0; i < nBuckets-1; i++)
				{

					BBox b0, b1;
					b0.Init();
					b1.Init();
					int count0 = 0, count1 = 0;
					for (int j = 0; j <= i; j++)
					{
						b0.Union(buckets[j].bounds);
						count0 += buckets[j].count;
					}
					for (int j = i+1; j < nBuckets; j++)
					{
						b1.Union(buckets[j].bounds);
						count1 += buckets[j].count;
					}
					float cost = eval.EvalCost(count0, SurfaceArea(b0), count1, SurfaceArea(b1), SurfaceArea(bbox));
#ifdef PROFILE_BVH
					if (nodeCount == 4)
					{
						/*GxWinSystem::Debug::Write(String(L"bbox0: "));
						GxWinSystem::Debug::Write(String(b0.Min.x) + L" ");
						GxWinSystem::Debug::Write(String(b0.Min.y) + L" ");
						GxWinSystem::Debug::Write(String(b0.Min.z) + L" ");
						GxWinSystem::Debug::Write(String(b0.Max.x) + L" ");
						GxWinSystem::Debug::Write(String(b0.Max.y) + L" ");
						GxWinSystem::Debug::Write(String(b0.Max.z) + L" ");
						GxWinSystem::Debug::WriteLine(L"");
						GxWinSystem::Debug::Write(String(L"bbox1: "));
						GxWinSystem::Debug::Write(String(b1.Min.x) + L" ");
						GxWinSystem::Debug::Write(String(b1.Min.y) + L" ");
						GxWinSystem::Debug::Write(String(b1.Min.z) + L" ");
						GxWinSystem::Debug::Write(String(b1.Max.x) + L" ");
						GxWinSystem::Debug::Write(String(b1.Max.y) + L" ");
						GxWinSystem::Debug::Write(String(b1.Max.z) + L" ");
						GxWinSystem::Debug::WriteLine(L"");*/
						GxWinSystem::Debug::WriteLine(String(L"cost:\t") + String(cost) + L"\t" + String(SurfaceArea(b0)) + L"\t" + String(SurfaceArea(b1)) + L"\t" + String(count0) + L"\t" + String(count1));
					}
#endif
					
					if (cost < minCost)
					{
						minCost = cost;
						minCostSplit = i;
#ifdef PROFILE_BVH
						minCostB1 = b0;
						minCostB2 = b1;
#endif
					}
				}
				if (elementCount > CostEvaluator::ElementsPerNode ||
					minCost < elementCount)
				{
					BuildData<T> *pmid = std::partition(elements,
						elements + elementCount,
						[&](const BuildData<T> &p)
						{
							int b = (int)(nBuckets * ((p.Center[dim] - centroidBounds.Min[dim]) /
								 (centroidBounds.Max[dim] - centroidBounds.Min[dim])));
							if (b == nBuckets) b = nBuckets-1;
							Assert(b >= 0 && b < nBuckets);
							return b <= minCostSplit;
						});
					node->Axis = dim;
	#ifdef PROFILE_BVH
					if (minCostSplit == 0 && depth<=32)
					{
						float endCost = eval.EvalCost(0, (bbox.Max[(dim+1)%3]-bbox.Min[(dim+1)%3])*(bbox.Max[(dim+2)%3]-bbox.Min[(dim+2)%3]), (int)elementCount, SurfaceArea(bbox), SurfaceArea(bbox));
						if (endCost < minCost)
						{
							int valAdd = 1;//1<<(32-depth);
							tree.Profile_Lock.Lock();
							tree.Profile_BadSplits += valAdd;
							tree.Profile_Lock.Unlock();
							node->IsBadNode = true;
						}
					}
					if (pmid == elements)
						pmid = elements + 1;
					else if (pmid == elements + elementCount)
						pmid = elements + elementCount-1;
	#endif
					job = BvhJob(node->Children, elements, pmid-elements);
					pushJob(node->Children+1, pmid, elements + elementCount - pmid);
				}
				else
				{
					node->AllocElements((int)elementCount);
					node->Bounds = bbox;
					for (int i = 0; i < (int)elementCount; i++)
					{
						node->Elements[i] = elements[i].Element;
					}
					elementListSize += elementCount;
					if (!stackPtr)
						break;
					else
						job = popJob();
					continue;
				}
			}
		}
		return rs;
	}

	template<typename T, typename CostEvaluator>
	void ConstructBvh(Bvh_Build<T> & tree, BuildData<T>* elements, int elementCount, CostEvaluator & eval)
	{
#ifdef PROFILE_BVH
		tree.Profile_BadSplits = 0;
#endif
		tree.Root = ConstructBvhNode<T,CostEvaluator>(tree, elements, elementCount, tree.ElementListSize, tree.NodeCount, eval, 0);
	}

	template<typename T, typename Tracer, bool pred>
	bool TraverseBvh(RayTraceSystem * system, const Tracer & tracer, DifferentialGeometry & rs, Bvh<T> & tree, const Ray & ray, const RayDifferential & rayDiff, float & t)
	{
#ifdef PROFILE_BVH
		int leavesTraversed = 0, nodesTraversed = 0, trianglesTested = 0, hits = 0, badNodes = 0;
#endif
		bool hit = false;
		Vec3 origin;
		Vec3::Scale(origin, ray.Direction, ray.tMin);
		Vec3::Add(origin, origin, ray.Origin);
		float tmax = ray.tMax;
		int dirIsNeg[3] = { ray.ReciprocalDirection.x < 0, ray.ReciprocalDirection.y < 0, ray.ReciprocalDirection.z < 0 };
		BvhNode* node = tree.Nodes.Buffer();
		int todoOffset = 0;
		BvhNode* todo[512];
		auto traceRay = ray;
		while (true)
		{

			float t1,t2;
			if (RayBBoxIntersection(node->Bounds, ray, t1, t2, dirIsNeg) && t1<traceRay.tMax)
			{
#ifdef PROFILE_BVH
				nodesTraversed++;
				if (node->IsBadNode)
					badNodes++;
#endif
				if (node->ElementCount > 0)
				{
#ifdef PROFILE_BVH
					leavesTraversed++;
#endif
					DifferentialGeometry inter;
					inter.Init();

					for (int i = node->ElementId; i < node->ElementId+node->ElementCount; i++)
					{
#ifdef PROFILE_BVH
						trianglesTested++;
#endif
						if (tracer.Trace(inter, tree.Elements[i], traceRay, rayDiff, tmax))
						{
							if (pred)
								return true;	
							if (tmax <= traceRay.tMax)
							{
								rs = inter;
								traceRay.tMax = tmax;
								t = tmax;
								hit = true;
#ifdef PROFILE_BVH
								hits++;
#endif
							}
						}
					}
					if (todoOffset == 0) break;
					node = todo[--todoOffset];
				}
				else
				{
					if (ray.Origin[node->Axis] > (node + 1)->Bounds.Max[node->Axis])
					{
						todo[todoOffset++] = node + 1;
						node = node + node->ChildOffset;
					}
					else
					{
						todo[todoOffset++] = node + node->ChildOffset;
						node = node + 1;
					}
				}
			}
			else
			{
				if (todoOffset == 0)
					break;
				node = todo[--todoOffset];
			}
		}
#ifdef PROFILE_BVH
		if (!pred)
			system->BvhProfiler.AddRecord(leavesTraversed, nodesTraversed, trianglesTested, hits, badNodes);
#endif
		return hit;
	}


	template <typename T, typename CreateLeafSpheresFunction>
	void ComputeBvhBoundingSpheres(Bvh<T> &tree, int depth, Array<Vec4, true> &buffer, CreateLeafSpheresFunction & cleaf)
	{
#ifdef DEBUG
		if (depth > 5)
			throw "Depth too large";
#endif
		int queue1[32], queue2[32];
		int *toTraverse = queue1, *nextQueue = queue2;
		int qPtr = 1, nPtr = 0;
		Vec3 position, range;
		Vec4 sphere;
		toTraverse[0] = 0;
		while (qPtr)
		{
			depth --;
			for (int i = 0; i<qPtr; i++)
			{
				int nodeId = toTraverse[i];
				if (tree.Nodes[nodeId].ElementCount)
				{
#ifdef DEBUG
					if (spheres >= bufferSize)
						throw "Output buffer overflow";
#endif
					cleaf(tree.Nodes[nodeId].Bounds, nodeId, depth);
				}
				else
				{
					if (depth == 0)
					{
#ifdef DEBUG
						if (spheres >= bufferSize)
							throw "Output buffer overflow";
#endif
						Vec3::Add(position, tree.Nodes[nodeId].Bounds.Min, tree.Nodes[nodeId].Bounds.Max);
						Vec3::Scale(position, position, 0.5f);
						sphere.x = position.x; sphere.y = position.y; sphere.z = position.z;
						Vec3::Subtract(range, tree.Nodes[nodeId].Bounds.Max, tree.Nodes[nodeId].Bounds.Min);
						sphere.w = range.Length() * 0.5f;
						buffer.Add(sphere);
					}
					else
					{
#ifdef DEBUG
						if (nPtr >= 32)
							throw "Buffer overflow";
#endif
						{
							nextQueue[nPtr] = nodeId+1;
							nPtr++;
						}
						int subId = nodeId + tree.Nodes[nodeId].ChildOffset;
						
						{
							nextQueue[nPtr] = subId;
							nPtr++;
						}
					}
					
				}
			}
			int * tmpQueue = nextQueue;
			nextQueue = toTraverse;
			toTraverse = tmpQueue;
			qPtr = nPtr;
			nPtr = 0;
		}

	}
}

#endif